module basicmodule
	use myfuncs
	use dotops
	use ML

	integer, parameter	:: ns=3000, nskip=50, nburnin=25000		! number of draws to save, number of posterior draws skipped, number of burnin draws
	logical, parameter	:: trialrun=.false.						! flag for Geweke test
	logical				:: worldfactor=.false.					! flag for factor informed by all countries (only OECD if false)
	logical				:: singlerun=.false. .or. trialrun		! flag for run with default values below only
	
	logical				:: poosflag=.true.						! flag for pseudo-out-of-sample calculation
	
	integer, parameter	:: n=merge(10,113,trialrun), maxh=100	
	integer, parameter	:: nclubs=merge(4,25,trialrun), nCoCs=merge(3,10,trialrun)
	integer, parameter	:: T0=118,periodfac0=merge(20,7,trialrun)
	integer				:: T=T0,periodfac=periodfac0
	integer				:: q0=T0/periodfac0
	integer				:: q=merge(7,T0/periodfac0+15,trialrun)
	integer, parameter	:: nbf=2, fhs(nbf)=[50,100], nh=maxh
	integer				:: Tmax=T0+maxh, npath=T0+nh

	integer, parameter	:: nf=25
	real, parameter		:: frhotab(nf)=0.5**(1.0/[(50+(i-1.0)*100.0/(nf-1),i=1,nf)])
	
	integer, parameter	:: nth0=100
	real				:: tscorrprop(3,nth0)
	integer, parameter	:: nth=merge(8,nth0,trialrun)
	integer, parameter	:: nrhogrid=25, ns2grid=25

	integer, parameter	:: ntdel=1
	integer, parameter	:: ssh(ntdel)=[50]		! horizon for vola calculations
	
	character(len=100)	:: dir="c:/dropbox/mystuff/SCC/out_play/"
	character(len=100)	:: cname="fMWW_15"
	
	real, allocatable	:: postmeanf(:)

! precomputed stuff
	real, allocatable	:: gammas(:,:)

	real, allocatable	:: Su(:,:,:), cholSu(:,:,:), Siu(:,:,:), Sfm(:,:,:), Sifm(:,:,:), cholfm(:,:,:), Sfa(:,:), cholfa(:,:), Sifa(:,:), G(:,:), Delta(:,:), Deltainv(:,:)
	real, allocatable	:: mfcstfm(:,:,:), cholfcstfm(:,:,:),mfcstfa(:,:), cholfcstfa(:,:), Xfcstf(:,:)
	real, allocatable	:: mfcstu(:,:,:), cholfcstu(:,:,:), Sfcstu(:,:,:)
	real, allocatable	:: detSu(:), Suss(:,:), Sutvo(:,:)
	real, allocatable	:: hlpriordist(:)
!	
	real				:: rhogrid(nrhogrid)=merge(0.8,0.95,trialrun)*[(real(i-1)/(nrhogrid-1),i=1,nrhogrid)]
!	real				:: Su(0:q,0:q,nth), cholSu(0:q,0:q,nth), Siu(0:q,0:q,nth)
!	real				:: detSu(nth), Suss(ntdel,nth), Sutvo(2,nth)
!	real				:: Sf(0:q,0:q,nf), Sif(0:q,0:q,nf), cholf(0:q,0:q,nf)
	real				:: detSfm(nf), sfmss(ntdel,nf), sfass(ntdel)
!	real				:: mfcstf(nh,0:q,nf), cholfcstf(nh,nh,nf)
!	real				:: mfcstu(nh,0:q,nth), cholfcstu(nh,nh,nth), Sfcstu(nh,nh,nth)
!	real				:: G(T,0:q)
	real, allocatable	:: SuAA(:,:,:,:),SuBB(:,:,:,:),SuAAS(:,:,:,:),SuBBS(:,:,:,:)
	real				:: s2grid(ns2grid)=-1, s2alpha(ns2grid)=-1
	
! state of sampler	
	real, allocatable	:: Us(:,:), Xs(:,:), f(:), CFs(:,:), CoCFs(:,:), ydummy(:)
	real, allocatable	:: fcstf(:), fcsts(:,:), fcstUs(:,:), fcstCFs(:,:),fcstCoCFs(:,:)
	real, allocatable	:: fm(:)
!	real				:: Us(0:q,n), Xs(0:q,n), f(0:q), CFs(0:q,nclubs), CoCFs(0:q,nCoCs), ydummy(0:q)
	real				:: sigma2_F(2),mt_F(2),muU,om2
	real				:: thdist(nth),CFthdist(nth),CoCthdist(nth)
	integer				:: cindu(n),cindf,cindCF(nclubs)
	real				:: rhos(n)
	integer				:: clubid(n)
	real				:: kappa2(n),clubs2(nclubs), CoCs2(nCoCs)
	integer				:: CoCid(nclubs),cindCoC(nCoCs)
	real				:: CoCrhos(nclubs)
	real				:: rhodist(nrhogrid),CoCrhodist(nrhogrid)	
	real				:: s2dist(ns2grid), CFs2dist(ns2grid), CoCs2dist(ns2grid)

!	real				:: fcstf(nh), fcsts(nh,n), fcstUs(nh,n)
!	real				:: fcstCFs(nh,nclubs),fcstCoCFs(nh,nCoCs)
	
	
! data	
	type tregion
		integer	:: qi
		real, allocatable	:: Y(:), w(:,:),AApAi(:,:)
	end type tregion	

	type(tregion)		:: region(n)
	real				:: pop(n),fweights(n)
	
! priors
	real, parameter		:: Deltavar=merge(2.0,0.01**2,trialrun)
	real, parameter		:: sigma2_Fanuprior=merge(5,1,trialrun), sigma2_Faprior=(0.03**2/2.198)*sigma2_Fanuprior
	integer, parameter	:: ns2fm=25
	real, parameter		:: s2fmtab0(ns2fm)=(0.01*[(0.1+(2-0.1)*(i-1)/(ns2fm-1),i=1,ns2fm)])**2
	real				:: s2fmtab(ns2fm)=s2fmtab0
	real, parameter		:: om2nuprior=merge(5,1,trialrun), om2prior=merge(0.50,1.0/2.198,trialrun)
	real, parameter		:: alphatot_prior=20, muprec_prior=merge(1.0,0.0,trialrun)
	
	character(len=:),allocatable	:: names(:)
	real, allocatable	:: tracestats(:,:)
	
end module

module prep
!DIR$ NOOPTIMIZE
	use basicmodule
	implicit none
	
	real		:: cutoff
	real, allocatable :: SRW(:,:), Xraw(:,:)
	real	, allocatable	:: Fw(:,:),ssv(:,:)
	contains
	
	subroutine sets2grid
		use chiin_int
		use chipr_int
		use chidf_int
		integer	:: i
		real	:: s2grid_e(0:ns2grid)
		s2grid_e=0.0
		do i=1,ns2grid
			s2grid(i)=(1.0/3)*(3.0)**(2*real(i-1)/(ns2grid-1))
			s2alpha(i)=1
		enddo
		s2alpha=s2alpha*alphatot_prior/sum(s2alpha)
		print *,"s2grid"
		call mdisp(s2grid)
		call mdisp(s2alpha)
	end subroutine
	
	function phifunc(j,s) result(val)
		integer		:: j
		real		:: s,val
		
		if(j==0) then
			val=1; return
		endif
		val=sqrt(2.0)*cos(Pi*s*j)
	end function
	
	subroutine initFw
		use evesf_int
		integer	:: i,s,j
		real		:: evals(q-1), evecs(T,q-1), m(T,T), Xrawfcst(0:Tmax-T,2)

		if(allocated(SRW)) then
			deallocate(SRW,Xraw,Fw,ssv)
			deallocate(Su,cholSu,Siu,Sfm,Sifm,cholfm,Sfa,Sifa,cholfa,G,Delta,Deltainv,Us,Xs,f,CFs,CoCFs,ydummy,mfcstfm,cholfcstfm,mfcstfa,cholfcstfa,Xfcstf)
			deallocate(mfcstu,cholfcstu,Sfcstu,fcstf,fcsts,fcstUs,fcstCFs,fcstCoCFs,fm)
			deallocate(detSu, Suss, Sutvo, hlpriordist)
		endif
		
		allocate(SRW(T,T), Xraw(T,2),Fw(Tmax,0:q+nh),ssv(tmax,ntdel))
		allocate(Su(0:q,0:q,nth), cholSu(0:q,0:q,nth), Siu(0:q,0:q,nth), Sfm(0:q,0:q,nf), Sifm(0:q,0:q,nf), cholfm(0:q,0:q,nf), Sfa(0:q,0:q))
		allocate(cholfa(0:q,0:q), Sifa(0:q,0:q), G(T,0:q),Delta(0:q,0:q),Deltainv(0:q,0:q))
		allocate(Us(0:q,n), Xs(0:q,n), f(0:q), CFs(0:q,nclubs), CoCFs(0:q,nCoCs), ydummy(0:q),fm(0:q))
		allocate(mfcstfm(nh,0:q,nf), cholfcstfm(nh,nh,nf),mfcstfa(nh,0:q), cholfcstfa(nh,nh),Xfcstf(nh,2))
		allocate(mfcstu(nh,0:q,nth), cholfcstu(nh,nh,nth), Sfcstu(nh,nh,nth))
		allocate(fcstf(nh), fcsts(nh,n), fcstUs(nh,n), fcstCFs(nh,nclubs),fcstCoCFs(nh,nCoCs))
		allocate(detSu(nth), Suss(ntdel,nth), Sutvo(2,nth), hlpriordist(nth))

		Xraw(:,1)=1.0/T
		Xrawfcst(:,1)=1.0/T
		Xraw(:,2)=[(i/real(T),i=1,T)]
		Xrawfcst(:,2)=[(i/real(T),i=T,Tmax)]
		Xrawfcst(:,2)=Xrawfcst(:,2)-sum(Xraw(:,2))/T
		Xraw(:,2)=Xraw(:,2)-sum(Xraw(:,2))/T
		Xrawfcst=matmul(Xrawfcst,invertpd(matmul(transpose(Xraw),Xraw)))
		do i=1,T
			do j=1,T
				SRW(i,j)=real(min(i,j))
			enddo
		enddo
		m=eye(T)-matmul(Xraw,matmul(invertpd(matmul(transpose(Xraw),Xraw)),transpose(Xraw)))
		call evesf(q-1,matmul(m,matmul(SRW,m)),.false.,evals,evecs)
		cutoff=0.9999*evals(q0-1)
		Fw=0
		Fw(1:T,0:1)=Xraw
		Fw(1:T,2:q)=evecs
		Fw(T,q+1:)=-1.0
		do s=1,maxh
			Fw(T+s,q+s)=1.0
			Xfcstf(s,:)=Xrawfcst(s,:)-Xrawfcst(0,:)
		enddo
		Delta=Deltavar*matmul(transpose(Fw(1:T,0:q)),Fw(1:T,0:q))
		call linds(Delta,Deltainv)
	end subroutine
	
	subroutine setSfromga(ga,S) 
		real	:: ga(0:),S(Tmax,Tmax)
		real	:: gax(-Tmax+1:Tmax-1)
		integer	:: k
		
		gax(0)=ga(0)
		do k=1,Tmax-1
			gax(k)=ga(k)
			gax(-k)=ga(k)
		enddo
		do k=0,Tmax-1
			S(:,k+1)=gax(-k:-k+Tmax-1)
		enddo
	end subroutine
	
	subroutine mkgammas
		use faure_int
		integer, parameter	:: kmax=800
		type (d_imsl_faure), pointer  :: fstate
		real		:: x(3)
		real		:: hl(2),r(2), corr(0:kmax,2)
		integer		:: i,j

		if(allocated(gammas)) deallocate(gammas)
		allocate(gammas(1:kmax+1,nth))
		call faure_init(3, fstate)
		do i=1,nth
			call faure_next(fstate,x)
			hl=25+775*x(1:2)**2
			r=2**(-1.0/hl)
			do j=0,kmax
				corr(j,:)=r**j
			enddo
			gammas(:,i)=matmul(corr,[x(3),1-x(3)])
			hlpriordist(i)=count(gammas(2:,i)>0.5)
			tscorrprop(:,i)=[r,x(3)]
        enddo
		call faure_free(fstate)
		call mdisp(hlpriordist)
		print *,sum(hlpriordist)/nth
!		stop
	end subroutine

	subroutine initSFs
		use faure_int
		integer	:: i,j,s1,s2
		real	:: Sraw(Tmax,Tmax), Sall(0:q+nh,0:q+nh), S(0:q,0:q)
		real	:: A(Tmax,Tmax),frho

		print *,"enter initSFs"
		
		ssv=0
		ssv(1,:)=-1
		do i=1,ntdel
			ssv(ssh(i)+1,i)=1
		enddo
		
		do j=1,tmax
			do i=1,tmax
				A(i,j)=boole(i>=j)
			enddo
		enddo

!$omp parallel do private(Sraw,Sall,S,frho)		
		do i=1,nf
			frho=frhotab(i)
			do s1=1,tmax
				do s2=1,tmax
					Sraw(s1,s2)=frho**abs(s1-s2)
				enddo
			enddo
			
			Sraw=matmul(A,matmul(Sraw,transpose(A)))
			sfmss(:,i)=diagonal(matmul(transpose(ssv),matmul(Sraw,ssv)))
			Sall=matmul(transpose(Fw),matmul(Sraw,Fw))
			S=Sall(0:q,0:q)
			detSfm(i)=detpd(S)
			call linds(S,Sifm(:,:,i))
			Sfm(:,:,i)=S
			cholfm(:,:,i)=choleski(S)
			mfcstfm(:,:,i)=matmul(Sall(q+1:,0:q),Sifm(:,:,i))
			cholfcstfm(:,:,i)=choleski(Sall(q+1:,q+1:)-matmul(mfcstfm(:,:,i),Sall(0:q,q+1:)))
		enddo
		Sraw=matmul(A,transpose(A))
		sfass=diagonal(matmul(transpose(ssv),matmul(Sraw,ssv)))
		Sall=matmul(transpose(Fw),matmul(Sraw,Fw))
		S=Sall(0:q,0:q)
		Sfa=S
		cholfa=choleski(S)
		call linds(S,Sifa)
		mfcstfa=matmul(Sall(q+1:,0:q),Sifa)
		cholfcstfa=choleski(Sall(q+1:,q+1:)-matmul(mfcstfa,Sall(0:q,q+1:)))
	end subroutine
	
	subroutine initSUs
		integer	:: i
		real	:: Sraw(Tmax,Tmax), Sall(0:q+nh,0:q+nh), S(0:q,0:q)
		
		if(allocated(SuAA)) deallocate(SuAA,SuBB,SuAAS,SuBBS,gammas)
		allocate(SuAA(0:q,0:q,nth,n),SuBB(0:q,0:q,nth,n),SuAAS(0:q,0:q,nth,n),SuBBS(0:q,0:q,nth,n))
		
		call mkgammas
		
		print *,"enter initSUs"
		
!$omp parallel do private(i,Sraw,Sall,S) 
		do i=1,nth
			call setSfromga(gammas(:,i),Sraw)
			Sall=matmul(transpose(Fw),matmul(Sraw,Fw))
			S=Sall(0:q,0:q)
			SU(:,:,i)=S
			detSU(i)=detpd(S)
			cholSu(:,:,i)=choleski(S)
			call linds(S,SiU(:,:,i))
			mfcstu(:,:,i)=matmul(Sall(q+1:,0:q),SiU(:,:,i))
			sfcstu(:,:,i)=Sall(q+1:,q+1:)-matmul(mfcstU(:,:,i),Sall(0:q,q+1:))
			cholfcstu(:,:,i)=choleski(sfcstu(:,:,i))
			Suss(:,i)=diagonal(matmul(transpose(ssv),matmul(Sraw,ssv)))
		enddo
	end subroutine
		
	subroutine getlfweights(sel,w)
		use evesf_int
		logical	:: sel(T)
		real, allocatable	::	w(:,:)
		real	:: m(T,T), nsel, X(T,2)
		real	:: evecs(T,q0-1),evals(q0-1)
		integer	:: i,qw
		
		nsel=count(sel)
		do i=1,2
			X(:,i)=Xraw(:,i)*rboole(sel)
		enddo
		m=diagmat(rboole(sel))-matmul(X,matmul(invertpd(matmul(transpose(X),X)),transpose(X)))
		call evesf(q0-1,matmul(m,matmul(SRW,m)),.false.,evals,evecs)
		qw=count(evals>cutoff)	
		if(allocated(w)) deallocate(w)
		allocate(w(T,0:qw+1))
		w(:,0:1)=X
		w(:,2:)=evecs(:,1:qw)
	end subroutine
	
	subroutine setregionwA(r,ir,sel)
		use rlse_int
		type(tregion)	:: r
		integer	:: ir
		logical	:: sel(T)
		real	:: AB(0:q,0:q), Sig(0:q,0:q), r2(0:q)
		real, allocatable	:: A(:,:),B(:,:)
		integer	:: i,s
		real	:: sst,sse
		
	print *,"enter setregion"
		do s=2,T-1
			if(.not. ( sel(s-1) .or. sel(s+1))) sel(s)=.false.
		enddo
		call getlfweights(sel,r%w) 
		r%qi=size(r%w,2)-1
		r2=0
		do i=0,r%qi
			call rlse(r%w(:,i),Fw(1:T,0:q),AB(i,:),intcep=0,sst=sst,sse=sse)
			r2(i)=1-sse/sst
		enddo
		print *,"r2 of projection regression"
		call mdisp(r2)
		AB(r%qi+1:q,:)=0
		do i=r%qi+1,q
			AB(i,i)=1.0
		enddo
		if(trialrun) then
			A=AB(:,0:r%qi)
			AB=transpose(AB)
			B=AB(:,r%qi+1:q)
		else
			AB=transpose(AB)
			A=AB(:,0:r%qi)
		endif
		r%AApAi=matmul(A,invertpd(matmul(transpose(A),A)))
!$omp parallel do private(Sig)		
		do i=1,nth
			Sig=Su(:,:,i)
			SuAA(:,:,i,ir)=matmul(matmul(Sig,A),matmul(invertpd(matmul(transpose(A),matmul(Sig,A))),transpose(A)))
			SuAAS(:,:,i,ir)=Sig-matmul(SuAA(:,:,i,ir),Sig)
			if(trialrun) then
				SuBB(:,:,i,ir)=matmul(matmul(Sig,B),matmul(invertpd(matmul(transpose(B),matmul(Sig,B))),transpose(B)))
				SuBBS(:,:,i,ir)=Sig-matmul(SuBB(:,:,i,ir),Sig)
			endif
		enddo
	end subroutine

	subroutine setG
		G=matmul(Fw(1:T,0:q),invertpd(matmul(transpose(Fw(1:T,0:q)),Fw(1:T,0:q))))
	end subroutine
	
	subroutine loadregions
		integer	:: i,ir
		real, allocatable	:: mdata(:,:),leveldata(:,:)
		real	:: unif(1)
		if(trialrun) then
			do ir=1,n
				associate(r=>region(ir))
					call rnun(unif)
					call setregionwA(r,ir,[(i>mod(ir*17,T),i=1,T)])
					allocate(r%Y(0:r%qi))					
					call rnnoa(r%Y)
				end associate
			enddo
			call rnnoa(fweights)
			fweights=fweights**2
			fweights(4:)=0
			fweights=fweights/sum(fweights)
		else
			names=loadstrings("C:\Dropbox\SCC\data\data_20191103\names.txt")
			mdata=loadmat("C:\Dropbox\SCC\data\data_20191103\pop_raw.csv")
			pop=sum(mdata(66:75,:),dim=1)/10
			print *,"pop"
			call mydispnames(names(1:n))
			call mdisp(pop)
			
			mdata=loadmat("C:\Dropbox\SCC\data\data_20191103\yp_raw.csv")
			if(poosflag) then
				mdata=mdata(1:T,:)
			else
				mdata=mdata(size(mdata,1)-T+1:,:)
			endif
			leveldata=mdata
			fweights=0
			do i=1,n
				associate(cname=>names(i))
				if(cname(4:4)=="x" .or. worldfactor) then
					fweights(i)=pop(i)
				endif
				end associate
			enddo
			fweights=fweights/sum(fweights)
			call mydispnames(names)
			call mdisp(fweights)
			block
				integer	:: i
				real		::  weff
				f=0;weff=0
				do i=1,n
					if(fweights(i)>0 .and. all(leveldata(:,i)==leveldata(:,i))) then
						F=f+fweights(i)*matmul(transpose(Fw(1:T,0:q)),log(leveldata(:,i)))
						weff=weff+fweights(i)
					endif
				enddo
				f=f/weff
				mt_f=f(0:1)

				call mdisp(f)
				call mdisp(matmul(G,f))
			end block
			do i=1,n
				associate(r=>region(i))
					print *,i, names(i)
					call setregionwA(r,i,leveldata(:,i)==leveldata(:,i))
					r%Y=matmul(transpose(r%w),log(merge(leveldata(:,i),1.0,leveldata(:,i)==leveldata(:,i))))
				end associate
			enddo
		endif
	end subroutine
end module
	
